#include "CorelSpy.h"
#include "../Share/Include/Common.h"

BOOL FormPath(LPTSTR szBuffer, LPTSTR szFolder, LPTSTR szFileName)
{
    if (!szBuffer || !szFolder || !szFileName)
    {
        return FALSE;;
    }

    // Init the buffer
    szBuffer[0] = _T('\0');
    _tcsncpy_s(szBuffer, MAX_PATH, szFolder, MAX_PATH);
    if (_T('\\') != szFolder[_tcslen(szFolder)-1])
    {
        _tcsncat_s(szBuffer, MAX_PATH, _T("\\"), MAX_PATH);
    }

    _tcsncat_s(szBuffer, MAX_PATH, szFileName, MAX_PATH);
    _tcsncat_s(szBuffer, MAX_PATH, _T("\0"), MAX_PATH);

    return TRUE;
}

HRESULT CorelDetourProcess()
{
    STARTUPINFO t_kStartupInfo;
    ZeroMemory(&t_kStartupInfo, sizeof(STARTUPINFO));
    PROCESS_INFORMATION t_kPI;
    ZeroMemory(&t_kPI, sizeof(PROCESS_INFORMATION));
    t_kStartupInfo.cb = sizeof(STARTUPINFO);

    // Init the buffer that will hold the current working folder of the application
    TCHAR t_szCurDir[MAX_PATH];
    DWORD t_dwCount = GetCurrentDirectory(MAX_PATH, t_szCurDir);
    if (!t_dwCount)
    {
        std::cerr << _T("Unable to get current working folder!\n");
        return GetLastError();
    }

	TCHAR t_szIniPath[MAX_PATH];
	TCHAR t_szTargetPath[MAX_PATH];
	TCHAR t_szCmdArgus[MAX_CMD_ARGU_LENGTH];
	TCHAR t_szDetouredName[MAX_PATH];
    TCHAR t_szDetouredPath[MAX_PATH];
	TCHAR t_szInjectName[MAX_PATH];
    TCHAR t_szInjectPath[MAX_PATH];


	ZeroMemory(t_szIniPath, sizeof(t_szIniPath));
	ZeroMemory(t_szTargetPath, sizeof(t_szTargetPath));
	ZeroMemory(t_szCmdArgus, sizeof(t_szCmdArgus));
	ZeroMemory(t_szDetouredName, sizeof(t_szDetouredName));
	ZeroMemory(t_szDetouredPath, sizeof(t_szDetouredPath));
	ZeroMemory(t_szInjectName, sizeof(t_szInjectName));
    ZeroMemory(t_szInjectPath, sizeof(t_szInjectPath));

	GetPrivateProfileString(_T(APP_SECTION), _T(APP_TARGET), _T(DEFAULT_APP), t_szTargetPath, MAX_PATH, _tcscat(_tcscpy(t_szIniPath, t_szCurDir), _T("\\")_T(INI_FILE_NAME)));
	GetPrivateProfileString(_T(APP_SECTION), _T(APP_CMDARGUS), _T(DEFAULT_CMDARGUS), t_szCmdArgus, MAX_CMD_ARGU_LENGTH, _T(INI_FILE_NAME));
	GetPrivateProfileString(_T(APP_SECTION), _T(APP_DETOURED), _T(DEFAULT_DETOURED), t_szDetouredName, MAX_PATH, _T(INI_FILE_NAME));
	GetPrivateProfileString(_T(APP_SECTION), _T(APP_DETOURS), _T(DEFAULT_INJECT), t_szInjectName, MAX_PATH, _T(INI_FILE_NAME));

    if (!FormPath(t_szDetouredPath, t_szCurDir, t_szDetouredName))
    {
        std::cerr << _T("Unable to form path to Detoured.dll!\n");
        return GetLastError();
    }

    if (!FormPath(t_szInjectPath, t_szCurDir, t_szInjectName))
    {
        std::cerr << _T("Unable to form path to CorelSpy.dll!\n");
        return GetLastError();
    }

    DWORD t_dwRet = SetEnvironmentVariable(ENVSTR_SERVER_PATH, t_szCurDir);
    if (t_dwRet == 0)
    {
        std::cerr << _T("Unable to set environment variable for Server path");
        return GetLastError();
    }

#ifdef UNICODE
	// Use short path workaround to make dll path can be placed in unicode path folder
	wchar_t t_szDetouredShort[MAX_PATH];
	wchar_t t_szInjectShort[MAX_PATH];
	char t_szDetouredPath2[MAX_PATH*2];
	char t_szInjectPath2[MAX_PATH*2];

	ZeroMemory(t_szDetouredShort, sizeof(t_szDetouredShort));
	ZeroMemory(t_szInjectShort, sizeof(t_szInjectShort));
	ZeroMemory(t_szDetouredPath2, sizeof(t_szDetouredPath2));
	ZeroMemory(t_szInjectPath2, sizeof(t_szInjectPath2));

	GetShortPathName(t_szDetouredPath, t_szDetouredShort, MAX_PATH);
	GetShortPathName(t_szInjectPath, t_szInjectShort, MAX_PATH);

	::WideCharToMultiByte(CP_ACP, 0, t_szDetouredShort, -1, t_szDetouredPath2, MAX_PATH*2, 0, 0);
	::WideCharToMultiByte(CP_ACP, 0, t_szInjectShort, -1, t_szInjectPath2, MAX_PATH*2, 0, 0);
#else
	char *t_szDetouredPath2 = t_szDetouredPath;
	char *t_szInjectPath2 = t_szInjectPath;
#endif

    BOOL t_bRet = DetourCreateProcessWithDll(
        t_szTargetPath,             // application name
        t_szCmdArgus,               // command line
        0,                          // process attribute
        0,                          // thread attribute
        0,                          // inherit handles
        CREATE_SUSPENDED|CREATE_DEFAULT_ERROR_MODE,  // creation flags
        0,                          // environment
        0,                          // current directory
        &t_kStartupInfo,            // startup info
        &t_kPI,                     // process information
        t_szDetouredPath2,          // detored dll full name
        t_szInjectPath2,            // inject dll full name
        0);                         // process creation callback

    if (!t_bRet)
    {
        std::cerr << _T("Cannot create process!\n");
        return GetLastError();
    }

    // Resume thread and wait on the process
    ResumeThread(t_kPI.hThread);

    WaitForSingleObject(t_kPI.hProcess, INFINITE);

    return S_OK;
}